import random 
import torch
import argparse
import numpy as np
import os
import pickle

from modules.SegModel import *
from driver.EDUSegmenter import *
from basic.Config import *
from driver.Dataloader import read_corpus

from TrainTest import scripts_evaluate, predict

from transformers import AutoModel, AutoTokenizer, AutoConfig, set_seed
from loguru import logger

if __name__ == '__main__':
    ### process id
    logger.info("Process ID {}, Process Parent ID {}".format(os.getpid(), os.getppid()))

    set_seed(666)

    ### gpu
    gpu = torch.cuda.is_available()
    print("GPU available: ", gpu)
    print("CuDNN: \n", torch.backends.cudnn.enabled)

    argparser = argparse.ArgumentParser()
    argparser.add_argument('--config_file', default='examples/default.cfg')
    argparser.add_argument('--model', default='BaseParser')
    argparser.add_argument('--thread', default=4, type=int, help='thread num')
    argparser.add_argument('--use-cuda', action='store_true', default=True)


    args, extra_args = argparser.parse_known_args()
    config = Configurable(args.config_file, extra_args)

    config.use_cuda = False
    if gpu and args.use_cuda: config.use_cuda = True

    logger.info(f"GPU using status:\t{config.use_cuda}")
    logger.info(f'Load pretrained plm: {config.plm_load_dir}')

    plm_model = AutoModel.from_pretrained(config.plm_load_dir)
    plm_config = AutoConfig.from_pretrained(config.plm_load_dir)
    plm_tokenizer = AutoTokenizer.from_pretrained(config.plm_load_dir, trust_remote_code=True)

    logger.info(f'Load pretrained plm ok')

    vocab = pickle.load(open(config.load_vocab_path, 'rb'))
    seg_model = SegModel(plm_model, plm_config, vocab)

    seg_model.output_layer.load_state_dict(torch.load(config.load_model_path)['dec'])
    if config.use_cuda:
        seg_model.cuda()

    segmenter = EDUSegmenter(seg_model, config)
    
    print("Test:")
    test_data = read_corpus(config.test_file)
    predict(test_data, segmenter, vocab, config, plm_tokenizer, config.test_file + '.out' )
    scripts_evaluate(config, config.test_file, config.test_file + '.out')